package org.example.using;

import org.apache.log4j.Level;
import org.apache.log4j.Logger;
import org.apache.spark.sql.*;

import static org.apache.spark.sql.functions.*;

import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.sql.expressions.Window;
import org.apache.spark.sql.expressions.WindowSpec;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.Metadata;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;
import org.example.bean.Employee;


import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;


/***
 * 创建一个类,用来测试 spark 2.3.0版本的 DataFrame的用法, DataSet API中窗口函数的使用.
 */
public class Spark05_SparkDataSetWindowUsing {

    public static void main(String[] args) throws IOException {
        // 初始化SparkSession
        // 引入log4j,消除控制台的警告
        Logger.getLogger("org").setLevel(Level.WARN);
        Logger.getLogger("akka").setLevel(Level.WARN);
        SparkConf conf = new SparkConf().setAppName("PantographModelFunction") // 设置应用程序名称
                .setMaster("local[*]") // 设置运行模式，这里是本地模式
                .set("spark.executor.memory", "6g") // 设置executor内存大小
                .set("spark.serializer", "org.apache.spark.serializer.KryoSerializer") // 设置序列化器
                .set("spark.driver.allowMultipleContexts", "true") // 允许多个SparkContext同时存在
                .set("spark.hadoop.validateOutputSpecs", "false") // 禁止输出路径已经存在时抛出异常
                .set("hive.mapred.supports.subdirectories", "true") // 开启Hive支持
                .set("mapreduce.input.fileinputformat.input.dir.recursive", "true") // 开启递归读取文件夹中的文件
                .set("spark.debug.maxToStringFields", "100").set("spark.sql.codegen.wholeStage", "false")
                .set("spark.sql.caseSensitive", "false"); // Spark会自动处理大小写不匹配的情况
        JavaSparkContext sc = new JavaSparkContext(conf);
        SparkSession spark = SparkSession.builder().config(conf).getOrCreate();


        // 定义一个集合,用来模拟 DataSet中的数据
        List<Employee> employeeList = new ArrayList<>();
        employeeList.add(new Employee("Alice", "HR", 5000));
        employeeList.add(new Employee("Bob", "HR", 6000));
        employeeList.add(new Employee("Charlie", "IT", 7000));
        employeeList.add(new Employee("David", "IT", 8000));
        employeeList.add(new Employee("Eva", "IT", 9000));
        // 这里是对窗口函数的使用示例
        // 创建一个 Employee类型的 DataSet.
        Dataset<Employee> employeeDS = spark.createDataset(employeeList, Encoders.bean(Employee.class));
        // 定义窗口函数中的 partition by 和 order by部分,也就是声明窗口规则.这里按照部门分组,按照 工资降序.
        // 这个版本 2.3.0 的 col() 等函数,还没有单独提出来作为一个方法,是 import static org.apache.spark.sql.functions 这个类的内部静态方法,在使用时需要注意导包的写法,同样下面使用到的 rank 也是.
        // 我们想要计算每个人的年龄的滑动平均值，窗口大小为当前行及其前后各一行。
        WindowSpec windowSpec = Window.partitionBy("department").orderBy(col("salary").desc()); // 定义窗口函数的规则
        Dataset<Row> result = employeeDS.withColumn("rank", rank().over(windowSpec)); // 窗口函数的使用
        result.show();


        // 关于滑动窗口的补充说明示例,这里演示基于行数的使用 `rowsBetween`
        List<Row> data2 = Arrays.asList(
                RowFactory.create("Alice", 30),
                RowFactory.create("Bob", 25),
                RowFactory.create("Carol", 35),
                RowFactory.create("Dave", 40),
                RowFactory.create("Eva", 20)
        );

        StructType schema2 = new StructType(new StructField[]{
                new StructField("name", DataTypes.StringType, false, Metadata.empty()),
                new StructField("age", DataTypes.IntegerType, false, Metadata.empty())
        });
        Dataset<Row> df2 = spark.createDataFrame(data2, schema2);
        WindowSpec windowSpec2 = Window.orderBy("age").rowsBetween(-1, 1);
        Dataset<Row> result2 = df2.withColumn("moving_avg", avg(col("age")).over(windowSpec2));
        result2.show();

        // 关于滑动窗口的补充说明示例,这里演示基于值范围定义窗口 `rangeBetween`
        List<Row> data3 = Arrays.asList(
                RowFactory.create("Alice", 3000),
                RowFactory.create("Bob", 5000),
                RowFactory.create("Carol", 8000),
                RowFactory.create("Dave", 6000),
                RowFactory.create("Eva", 4000)
        );

        StructType schema3 = new StructType(new StructField[]{
                new StructField("name", DataTypes.StringType, false, Metadata.empty()),
                new StructField("salary", DataTypes.IntegerType, false, Metadata.empty())
        });
        Dataset<Row> df3 = spark.createDataFrame(data3, schema3);
        // 在这个示例中，rangeBetween 根据工资范围进行了计算，结果显示了每个人工资在 ±2000 范围内的平均值。
        WindowSpec windowSpec3 = Window.orderBy("salary").rangeBetween(-2000, 2000);
        Dataset<Row> result3 = df3.withColumn("moving_avg", avg(col("salary")).over(windowSpec3));
        result3.show();

        // 在使用窗口时,使用
    }
}
